package chagine.core.restful.mvc;


import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.PrintWriter;

/**
 * 杭州乾瑾科技
 *
 * @author 王彬安
 * @create 2018/7/14 15:25
 */
public class WrapperedHttpServletResponse extends HttpServletResponseWrapper {
    private ByteArrayOutputStream buffer;
    private ServletOutputStream out;
    private PrintWriter writer;

    public WrapperedHttpServletResponse(HttpServletResponse resp) {
        super(resp);
    }

    private ByteArrayOutputStream getBuffer() {
        if (buffer == null) {
            // 真正存储数据的流
            buffer = new ByteArrayOutputStream(128);
        }
        return buffer;
    }

    /**
     * 重载父类获取outputstream的方法
     */
    @Override
    public ServletOutputStream getOutputStream() {
        if (out == null) {
            out = new WapperedOutputStream(getBuffer());
        }
        return out;
    }

    /**
     * 重载父类获取writer的方法
     */
    @Override
    public PrintWriter getWriter() {
        if (writer == null) {
            writer = new PrintWriter(getOutputStream());
        }
        return writer;
    }

    /**
     * 重载父类获取flushBuffer的方法
     */
    @Override
    public void flushBuffer() throws IOException {
        if (out != null) {
            out.flush();
        }
        if (writer != null) {
            writer.flush();
        }
    }

    @Override
    public void reset() {
        if (buffer != null) {
            buffer.reset();
        }
    }

    /**
     * 将out、writer中的数据强制输出到WapperedResponse的buffer里面，否则取不到数据
     */
    public byte[] getResponseData() throws IOException {
        if (buffer != null) {
            flushBuffer();
            return buffer.toByteArray();
        }
        return new byte[0];
    }

    public void writeToRealResponse() throws IOException {
        if (buffer != null) {
            buffer.writeTo(this.getResponse().getOutputStream());
        }
    }

    /**
     * 内部类，对ServletOutputStream进行包装
     */
    private static class WapperedOutputStream extends ServletOutputStream {
        private ByteArrayOutputStream bos;

        public WapperedOutputStream(ByteArrayOutputStream buffer) {
            super();
            bos = buffer;
        }

        @Override
        public void write(int b) {
            bos.write(b);
        }

        @Override
        public void write(byte[] b) {
            bos.write(b, 0, b.length);
        }

        @Override
        public void flush() throws IOException {
            bos.flush();
        }

        @Override
        public boolean isReady() {
            return true;
        }

        @Override
        public void setWriteListener(WriteListener writeListener) {

        }
    }
}
